#    Copyright 2011 Ragnvald Fuglestein
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.

from json import dumps as jsondumps
from json import loads as jsonloads
from pickle import dump as pickledump
from pickle import load as pickleload
from base64 import b64encode
from shutil import rmtree
from pyinotify import IN_CLOSE_WRITE
from pyinotify import IN_CREATE
from pyinotify import IN_MOVED_FROM
from pyinotify import IN_MOVED_TO
from pyinotify import IN_DELETE
from pyinotify import WatchManager
from pyinotify import Notifier
from os import path
from os import umask
from os import makedirs
from os import link
from os import remove
from os import stat
from urllib.request import Request
from urllib.request import urlopen
from urllib import error as urlerror
from fnmatch import fnmatch
from shlex import split

WatchDirs={}
LinkDict={}
GlobalConfig={'config':[]}
RpcHeaders={}
watchmanager=""
WatchFlags=IN_CLOSE_WRITE|IN_CREATE|IN_MOVED_FROM|IN_MOVED_TO|IN_DELETE
def Match(FileName,Filter):
   for FilterType in Filter:
      if fnmatch(FileName.lower(),FilterType.lower()):
         return True
   return False

def RpcReq(method,arguments,Retry=True):
   global RpcHeaders, RpcAddr
   try:
      Req = Request('http://127.0.0.1:9091/transmission/rpc',jsondumps({'method': method,'arguments': arguments}),RpcHeaders)
      try:
         Response=urlopen(Req)
      except urlerror.HTTPError as e:
         if e.code == 409:
            RpcHeaders['x-transmission-session-id']=e.headers['x-transmission-session-id']
            if Retry:
               return RpcReq(method,arguments,False)
         print(e)
         raise
      except urlerror.URLError as e:
         print(e.reason,": Is transmission-daemon running?")
         return False
   except urlerror.URLError as e:
      print(e)
      raise
   try:
      Reply=jsonloads(Response.read().decode('utf-8'))
   except TypeError as e:
      print("TypeError:",e,"\n",Reply)
      raise
   if 'result' not in Reply:
      print("Unknown error:",Reply)
      return False
   if Reply['result'] != 'success':
      print("No success:",Reply)
      return False
   if method == 'torrent-add':
      return Reply['arguments']['torrent-added']
   if method == 'torrent-get':
      return Reply['arguments']['torrents'][0]['files']
   if Reply['result'] == 'success':
      return True
   else:
      print("Weird error:",Reply)
   return False

def AddTorrent(TorrentFile,Config):
   torrentpathname=TorrentFile.pathname
   try:
      with open(torrentpathname, 'rb') as f:
         metainfo = b64encode(f.read()).decode('utf-8')
   except IOError as e:
      print(e)
      return False
   Torrent=RpcReq('torrent-add',{'metainfo': metainfo,'paused': False})
   if not Torrent:
      print("Error adding torrent:",torrentpathname)
      return False
   DownloadDir=Config['maindir']+Config['torrentdir'].replace("[:HASH:]",Torrent['hashString']).replace("[:NAME:]", 
   ''.join(c for c in Torrent['name'] if c in "-_.()[],'*&+=%% 0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"))
   if path.isdir(DownloadDir):
      print("Error: DownloadDir already exist:",DownloadDir)
      RpcReq('torrent-remove',{'ids': [Torrent['id']], 'delete-local-data': True})
      return False
   
   mask=umask(0o000)
   makedirs(DownloadDir)
   umask(mask)
   if not RpcReq('torrent-set-location',{'ids': Torrent['id'], 'location': DownloadDir, 'move': True}):
      print("Error moving torrent:",torrentpathname)
      RpcReq('torrent-remove',{'ids': [Torrent['id']], 'delete-local-data': True})
      rmtree(DownloadDir)
      return False

   FileList=RpcReq('torrent-get',{'ids': Torrent['id'], 'fields': ['files']})
   if not FileList:
      print("Error getting filelist")
      RpcReq('torrent-remove',{'ids': [Torrent['id']], 'delete-local-data': True})
      rmtree(TorrentDir)
      return False
   
   for Item in Config['config']:
      try:
         Option, Value=Item.split(" ",1)
      except ValueError:
         if Item == "denyall":
            for File in FileList:
               File['get']=False
         elif Item == "allowall":
            for File in FileList:
               File['get']=True
         elif Item == "removetorrent":
            remove(torrentpathname)
         elif Item == "start":
            if not RpcReq('torrent-start',{'ids': [Torrent['id']]}):
               print("Warning: torrent-start failed")
         else:
            print("Unknown Option:",Item)
      else:
         if Option == "allow":
            for File in FileList:
               FileName=File['name'].rsplit('/',1)
               if len(FileName)==2:
                  FileName.pop(0)
               if Match(FileName[0],split(Value)):
                  File['get']=True
         elif Option == "deny":
            for File in FileList:
               FileName=File['name'].rsplit('/',1)
               if len(FileName)==2:
                  FileName.pop(0)
               if Match(FileName[0],split(Value)):
                  File['get']=False
         elif Option == "ratio":
            if not RpcReq('torrent-set',{'ids': [Torrent['id']], 'seedRatioLimit': Value}):
               print("Warning: Invalid ratio:",Value)
         else:
            print("Unknown Options:%s %s"%(Option,Value))
   UnwantedFiles=[FileID for FileID,File in enumerate(FileList) if not File['get']]
   if UnwantedFiles!=[]:
      if not RpcReq('torrent-set',{'ids': [Torrent['id']], 'files-unwanted': UnwantedFiles}):
         print("Error changing wanted files:",FileList)
         RpcReq('torrent-remove',{'ids': [Torrent['id']], 'delete-local-data': True})
         rmtree(DownloadDir)
         return False
   global watchmanager, LinkDict, WatchFlags
   watchmanager.add_watch(DownloadDir,WatchFlags,proc_fun=FileEvent,rec=True,auto_add=True)
   LinkFiles=LinkDict[DownloadDir]={}
   
   for File in FileList:
      if File['get']:
         if 'basename' in Config or File['name'].find("/")==-1:
            LinkFiles['%s/%s'%(DownloadDir,File['name'])]="%s/%s" % (TorrentFile.path,path.split(File['name'])[1])
         else:
            LinkFiles['%s/%s'%(DownloadDir,File['name'])]="%s/%s" % (TorrentFile.path,File['name'])
   
def CreateLink(LinkFrom,LinkTo):
   Dir,File = LinkTo.rsplit("/",1)
   if not path.isdir(Dir):
      makedirs(Dir)
   if path.isfile(LinkTo):
      if stat(LinkFrom).st_ino!=stat(LinkTo).st_ino:
         print("Error: %s exists but not same inode"%(LinkTo))
         return False
   else:
      link(LinkFrom,LinkTo)
      if __debug__:
         print("link %s"%(LinkTo))
   return True
   
def FileEvent(File):
   global ProcessedFiles,WatchDirs,LinkDict
   PathName=File.pathname
   MaskName=File.maskname
   if not File.dir:
      if MaskName=="IN_CLOSE_WRITE" or MaskName=="IN_MOVED_TO":
         if File.name.endswith(".torrent"):
            for WatchedDir in WatchDirs:
               if File.path.startswith(WatchedDir):
                  return AddTorrent(File,WatchDirs[WatchedDir])

         for LinkDirs in LinkDict:
            LinkFiles=LinkDict[LinkDirs]
            if PathName in LinkFiles:
               if CreateLink(PathName,LinkFiles[PathName]):
                  del LinkFiles[PathName]
               if LinkFiles=={}:
                  if __debug__:
                     print("UnwatchDir %s"%(LinkDirs))
                  del LinkDict[LinkDirs]
               return
   else:
      if MaskName=="IN_DELETE_SELF" or MaskName=="IN_MOVE_SELF":
         print("IN_DELETE_SELF or IN_MOVE_SELF:%s"%(File))
         if PathName in LinkDict:
            if __debug__:
               print("UnwatchDir %s"%(File.path))
            del LinkDict[PathName]
      if MaskName=="IN_DELETE":
         print("IN_DELETE:%s"%(File))

def CheckConfig():
   global GlobalConfig,WatchDirs
   for W in WatchDirs.keys():
      if not path.isdir(W):
         print("SanityCheck error: %s does not exist"%(W))
         return False
      WatchDir=WatchDirs[W]
      if not 'rpcaddr' in WatchDir and not 'rpcaddr' in GlobalConfig:
         WatchDir['rpcaddr'] = "127.0.0.1:9091"
      if not 'maindir' in WatchDir:
         if not 'maindir' in GlobalConfig:
            print("SanityCheck error: No maindir!")
            return False
         WatchDir['maindir'] = GlobalConfig['maindir']
      if not path.isdir(WatchDir['maindir']):
         print("SanityCheck error: %s does not exist"%(WatchDir['maindir']))
         return False
      if not 'torrentdir' in WatchDir:
         if not 'torrentdir' in GlobalConfig:
            print("SanityCheck error: No torrentdir specified for %s"%(W))
            return False
         WatchDir['torrentdir'] = GlobalConfig['torrentdir']
      WatchDir['config']=["allowall"]+GlobalConfig['config']+WatchDir['config']
   if 'linkpickle' not in GlobalConfig:
      if 'maindir' in GlobalConfig:
         GlobalConfig['linkpickle']=GlobalConfig['maindir']+"/Pickle"
      else:
         print("SanityCheck error: No linkpickle or maindir defined globally")
         return False
   return True

def ReadPickle(LinkPickle,LinkDict):
   try:
      with open(LinkPickle,mode='rb') as f:
         LinkDict=pickleload(f)
   except (ValueError,IOError) as e:
      print("Warning: Creating new pickle. Error=%s"%(e))
      LinkDict={}
   else:
      UnwatchDir=[]
      for WDir in LinkDict:
         LinkDirs=LinkDict[WDir]
         if path.isdir(WDir):
            UnwatchFile=[]
            for LinkFrom in LinkDirs:
               LinkTo=LinkDirs[LinkFrom]
               if path.isfile(LinkFrom):
                  if CreateLink(LinkFrom,LinkTo):
                     UnwatchFile.append(LinkFrom)
            for PathName in UnwatchFile:
               del LinkDirs[PathName]
            if LinkDirs=={}:
               if __debug__:
                  print("UnwatchDir %s"%(WDir))
               UnwatchDir.append(WDir)
         else:
            print("Ignoring files in removed folder: %s"%(WDir))
            UnwatchDir.append(WDir)
      for PathName in UnwatchDir:
         del LinkDict[PathName]

def ReadConfig(ConfigFile):
   global GlobalConfig
   WatchDir=GlobalConfig
   try:
      with open(ConfigFile,encoding='utf-8',mode='r') as f:
         for Line in f:
            Option=Line.rstrip()
            if Option and not Option.startswith("#"):
               if path.isdir(Option):
                  WatchDir=WatchDirs[Option]={'config':[]}
               else:
                  if Option.startswith("linkpickle") and WatchDir==GlobalConfig:
                     WatchDir['linkpickle']=Option.split(' ',1)[1]
                  elif Option.startswith("rpcaddr"):
                     WatchDir['rpcaddr']=Option.split(' ',1)[1]
                  elif Option.startswith("maindir"):
                     WatchDir['maindir']=Option.split(' ',1)[1]
                  elif Option.startswith("torrentdir"):
                     WatchDir['torrentdir']=Option.split(' ',1)[1]
                  elif Option == "basename":
                     WatchDir['basename']=True
                  else:
                     WatchDir['config'].append(Option)
   except IOError as e:
      print("Error reading config %s\n%s"%(ConfigFile,e))
      return False
   return True
   
def main():
   global WatchDirs,GlobalConfig,LinkDict,watchmanager,WatchFlags
   
   if not ReadConfig("/etc/transmission-linker.conf"):
      return False
   if not CheckConfig():
      return False
   LinkPickle=GlobalConfig['linkpickle']
   if path.isfile(LinkPickle):
      ReadPickle(LinkPickle,LinkDict)
      
   SelectedWatchDirs=list(WatchDirs.keys())+list(LinkDict.keys())
   for WatchDir in [WatchDirs[W] for W in WatchDirs]:
      if not WatchDir['maindir'] in SelectedWatchDirs:
         SelectedWatchDirs.append(WatchDir['maindir'])
   watchmanager = WatchManager()
   notifier = Notifier(watchmanager)
   watchmanager.add_watch(SelectedWatchDirs,WatchFlags,proc_fun=FileEvent,rec=True,auto_add=True)
   while True:
      try:
         notifier.process_events()
         if notifier.check_events():
            notifier.read_events()
      except:
         notifier.stop()
         with open(LinkPickle,mode='wb') as f:
            pickledump(LinkDict,f)
         raise

if __name__ == "__main__":
   main()
